#include <benchmark/benchmark.h>

#include <memory>
#include <random>
#include <string>

#include "data_dumper.h"

static void BM_DataDumperReport_Buffer(benchmark::State& state) {
  uint32_t ring_buffer_size = state.range(0);
  uint32_t report_count = state.range(1);

  // 为每个测试用例创建唯一的路径
  std::string test_path = "./tmp/benchmark_test_" +
                          std::to_string(ring_buffer_size) + "_" +
                          std::to_string(report_count);

  // DataDumper dumper;
  // dumper.Init(test_path, ring_buffer_size);
  // dumper.Start();
  DataDumper::GetInstance().Init(test_path, ring_buffer_size);
  DataDumper::GetInstance().Start();

  for (auto _ : state) {
    state.PauseTiming(); // 暂停计时器
    std::vector<std::unique_ptr<OpRangeData>> data_vec;
    for (uint32_t i = 0; i < report_count; ++i) {
      auto report_data = std::make_unique<OpRangeData>(
          0, 0, 0, 0, 0, 0, 0, false, "test_tag");
    
      report_data->input_dtypes.push_back("test_dtype1");
      report_data->input_dtypes.push_back("test_dtype2");
      report_data->input_dtypes.push_back("test_dtype3");

      report_data->input_shapes.push_back({1, 2, 3});
      report_data->input_shapes.push_back({4, 5, 6});
      report_data->input_shapes.push_back({7, 8, 9});

      report_data->stack.push_back("test_stack1.py/test_func1/test_line1");
      report_data->stack.push_back("test_stack2.py/test_func2/test_line2");
      report_data->stack.push_back("test_stack3.py/test_func3/test_line3");

      report_data->module_hierarchy.push_back("test_module1");
      report_data->module_hierarchy.push_back("test_module2");
      report_data->module_hierarchy.push_back("test_module3");

      data_vec.push_back(std::move(report_data));
    }
    state.ResumeTiming(); // 恢复计时器

    for (auto& report_data : data_vec) {
      DataDumper::GetInstance().Report(std::move(report_data));
    }
  }

  DataDumper::GetInstance().Stop();
}

BENCHMARK(BM_DataDumperReport_Buffer)
    ->Ranges({{128, 8192}, {10000, 100000}})
    ->Unit(benchmark::kMicrosecond);

BENCHMARK_MAIN();
